//
//  ToneMapperController.m
//  ToneMapper
//
//  Created by Andrey on 03/02/2009.
//  Copyright 2009 Karma Software. All rights reserved.
//

#import "ToneMapperController.h"
#import "ToneMapperView.h"
#import "HistogramView.h"
#import "ChromaticityHistogramView.h"
#import "KImageSegmentator.h"
#import "KSimpleFloatArray.h"
#import "KSegmentationMap.h"
#import "KCurve.h"
#import "CurvesView.h"
#import "KHistogram.h"
#import "KCumulativeHistogram.h"
#import "K2DHistrogram.h"

CGFloat GetEqualizedColorComponent(NSUInteger* componentHistogramData, 
								   NSUInteger histogramIntervalCount, 
								   CGFloat component)
{
	CGFloat histogramSpan = 
	componentHistogramData[histogramIntervalCount - 1] - componentHistogramData[0];
	
	NSUInteger interval = (NSUInteger)(histogramIntervalCount * component);
	if (interval == histogramIntervalCount)
		interval = histogramIntervalCount - 1;
	
	return (componentHistogramData[interval] - componentHistogramData[0]) / histogramSpan;
}

@implementation ToneMapperController

@synthesize currentImage;
@synthesize displayBuffer;
@synthesize segmentationMap;
@synthesize chromaticitySegmentationMap;
@synthesize channelForSegmentation;
@synthesize useChromaticityForSegmentation;
@synthesize displayedHistogramChannel;
@synthesize displayedChromaticityHistogramChannel;
@synthesize displayCumulativeHistogram;

- (id) init
{
	if (self = [super init])
	{
		currentImage = nil;
		displayBuffer = nil;
		displayedHistogramChannel = BrightnessImageChannel;
		displayCumulativeHistogram = NO;
		useChromaticityForSegmentation = NO;
	}
	
	return self;
}

- (IBAction) openDocument: sender  
{
	[[NSOpenPanel openPanel] 
	 beginSheetForDirectory:nil
	 file:nil
	 modalForWindow:window 
	 modalDelegate:self
	 didEndSelector:@selector(openPanelDidEnd:returnCode:contextInfo:)
	 contextInfo:NULL];  
}

- (IBAction) saveDocument: (id) sender
{
	[[NSSavePanel savePanel]
	 beginSheetForDirectory:nil
	 file:nil
	 modalForWindow:window 
	 modalDelegate:self
	 didEndSelector:@selector(savePanelDidEnd:returnCode:contextInfo:)
	 contextInfo:NULL];  
}

- (void) openPanelDidEnd:(NSOpenPanel *)sheet returnCode:(int)returnCode contextInfo:(void*)contextInfo
{
	if (returnCode == NSOKButton)
	{
		if (currentImage != nil)
			[currentImage release];
		
		currentImage = [[[KImage alloc] initWithFile:[sheet filename]] retain];
		displayBuffer = currentImage;
		
		NSSize currentImageSize = NSMakeSize([currentImage width], [currentImage height]);
		NSPoint newWindowOrigin = 
		NSMakePoint([window frame].origin.x - (currentImageSize.width - [window frame].size.width),
					[window frame].origin.y - (currentImageSize.height - [window frame].size.height));
		NSRect newWindowFrame = 
		NSMakeRect(newWindowOrigin.x, newWindowOrigin.y, currentImageSize.width, currentImageSize.height);
		
		[window setFrame: newWindowFrame display: YES animate: YES];
		[window setMaxSize: currentImageSize];
		
		[mainView setFrameSize: currentImageSize];
		[mainView setBoundsSize: currentImageSize];
		
		[histogramView setNeedsDisplay:YES];
		[chromaticityHistogramView setNeedsDisplay:YES];
	//	[self showChromaticityHistogramWindow:self];
	//	[self showHistogramWindow:self];
		[mainView setNeedsDisplay: YES];
	}
}

- (void)savePanelDidEnd:(NSSavePanel *)sheet returnCode:(int)returnCode contextInfo:(void*)contextInfo
{
	if (returnCode == NSOKButton)
	{
		NSData* tiffData = [displayBuffer.bitmap TIFFRepresentation];
		[tiffData writeToFile:[sheet filename] atomically:NO];
	}
}

- (void) backupCurrentImage
{
	NSData* tiffData = [currentImage.bitmap TIFFRepresentation];
	NSString* backupFilename = 
	[NSTemporaryDirectory() stringByAppendingString:@"CurrentImageBackup.tiff"];
	
	[tiffData writeToFile:backupFilename atomically:NO];
}

- (IBAction) restoreCurrentImageFromBackup: (id) sender
{
	NSString* backupFilename = 
	[NSTemporaryDirectory() stringByAppendingString:@"CurrentImageBackup.tiff"];
	[currentImage restoreFromBackup:backupFilename];
	
	[currentImage.histogram computeForImage:currentImage];
	[currentImage.cumulativeHistogram computeForImage:currentImage];
	[self updateHistogramStatistics];
	
	displayBuffer = currentImage;
	[histogramView setNeedsDisplay:YES];
	[chromaticityHistogramView setNeedsDisplay:YES];
	[mainView setNeedsDisplay:YES];
}

- (IBAction) showHistogramWindow: (id) sender
{	
	[histogramChannelComboBox selectItemAtIndex:0];
	displayedHistogramChannel = [histogramChannelComboBox indexOfSelectedItem];
	
	[self updateHistogramStatistics];
	
	[histogramView setNeedsDisplay:YES];
	[histogramWindow makeKeyAndOrderFront:histogramWindow];
}

- (IBAction) changeDisplayedHistogramChannel: (id) sender
{
	displayedHistogramChannel = [histogramChannelComboBox indexOfSelectedItem];
	
	[self updateHistogramStatistics];
	
	[histogramView setNeedsDisplay:YES];
}

- (void) updateHistogramStatistics
{
	if (currentImage != nil)
	{
		CGFloat channelMean = 
		currentImage.histogram.channelHistograms[displayedHistogramChannel].mean;
		[meanLabel setFloatValue:channelMean];
		
		CGFloat channelVariance = 
		currentImage.histogram.channelHistograms[displayedHistogramChannel].variance;
		[varianceLabel setFloatValue:channelVariance];
	}
}

- (IBAction) changeHistogramType: (id) sender
{
	if (displayCumulativeHistogram == NO)
		displayCumulativeHistogram = YES;
	else
		displayCumulativeHistogram = NO;
	
	[histogramView setNeedsDisplay:YES];
}

- (void) trackMouseDragging: (NSPoint)mouseLocation
{
	mouseLocation.y = currentImage.height - mouseLocation.y;
	KRedundantColorVector colorVector = 
	[[displayBuffer pixelAtX:mouseLocation.x y:mouseLocation.y] redundantColorVector];
	
	CGFloat component = colorVector.data[displayedHistogramChannel];
	
	NSUInteger highlightedInterval = 
	(NSUInteger)(component * currentImage.histogram.intervalCount);
	
	if (highlightedInterval == currentImage.histogram.intervalCount)
		highlightedInterval == currentImage.histogram.intervalCount - 1;
	
	histogramView.highlightedInterval = highlightedInterval;
	[histogramView setNeedsDisplay:YES];
	
	// Chromaticity histogram tracking
	chromaticityHistogramView.highlightedCell = 
	NSMakePoint(256 * colorVector.data[AImageChannel], 
				(1 - colorVector.data[BImageChannel]) * 256);
	[chromaticityHistogramView setNeedsDisplay:YES];
}

- (IBAction) showChannelsWindow: (id) sender
{
	[channelComboBox selectItemAtIndex:0];
	[channelsWindow orderFront:channelsWindow];
}

- (IBAction) changeChannel: (id) sender
{
	KImage* channelImage = [[KImage alloc] initWithWidth:currentImage.width andHeight:currentImage.height];
	
	ForEachImagePixel(currentImage)
	{
		KPixel* currentPixel = [currentImage pixelAtX:x y:y];
		KRedundantColorVector colorVector = [currentPixel redundantColorVector];
		
		CGFloat channelValue = 
		colorVector.data[[histogramChannelComboBox indexOfSelectedItem]];
		
		KPixel* newCurrentPixel = 
		[[KPixel alloc] initWithRed:channelValue green:channelValue andBlue:channelValue];
		
		[channelImage setPixel:newCurrentPixel atX:x y:y];
	}
	
	displayBuffer = channelImage;
	[mainView setNeedsDisplay:YES];
}

- (IBAction) showCurvesWindow: (id)sender
{
	[curvesWindow makeKeyAndOrderFront:curvesWindow];
}

- (IBAction) applyCurves: (id)sender
{
	[curvesWindow orderOut:self];
	[self backupCurrentImage];
	
	ForEachImagePixel(currentImage)
	{
		KPixel* currentPixel = [currentImage pixelAtX:x y:y];
			
		CGFloat brightness = [curvesView.curve valueAt:[currentPixel brightness]];
		CGFloat hue = [currentPixel hue];
		CGFloat saturation = [currentPixel saturation];
		
		KPixel* newCurrentPixel = 
		[[KPixel alloc] initWithHue:hue saturation:saturation andBrightness:brightness];
		
		[currentImage setPixel:newCurrentPixel atX:x y:y];
	}
	
	[currentImage.histogram computeForImage:currentImage];
	[currentImage.cumulativeHistogram computeForImage:currentImage];
	[self updateHistogramStatistics];
	[histogramView setNeedsDisplay:YES];
	[chromaticityHistogramView setNeedsDisplay:YES];
	[mainView setNeedsDisplay:YES];
}

- (IBAction) recalibrateBrightness: (id) sender
{
	[mainView setNeedsDisplay:YES];
}

- (void) equalizeHueHistogram
{
	KCumulativeHistogram* cumulativeHistogram = currentImage.cumulativeHistogram;
	
	ForEachImagePixel(currentImage)
	{
		KPixel* currentPixel = [currentImage pixelAtX:x y:y];
			
		CGFloat brightness = [currentPixel brightness];
		CGFloat hue = [currentPixel hue];
		CGFloat saturation = [currentPixel saturation];
			
		hue = 
		GetEqualizedColorComponent(cumulativeHistogram.channelHistograms[HueImageChannel].data,
								   cumulativeHistogram.intervalCount, hue);
		
		
		KPixel* newCurrentPixel = 
		[[KPixel alloc] initWithHue:hue saturation:saturation andBrightness:brightness];
		
		[currentImage setPixel:newCurrentPixel atX:x y:y];
	}
}

- (void) equalizeRGBHistogram
{
	KCumulativeHistogram* cumulativeHistogram = currentImage.cumulativeHistogram;
	
	ForEachImagePixel(currentImage)
	{
		KPixel* currentPixel = [currentImage pixelAtX:x y:y];
		CGFloat r = [currentPixel red];
		CGFloat g = [currentPixel green];
		CGFloat b = [currentPixel blue];
		
		r = GetEqualizedColorComponent(cumulativeHistogram.channelHistograms[RedImageChannel].data, 
									   cumulativeHistogram.intervalCount, r);

		g = GetEqualizedColorComponent(cumulativeHistogram.channelHistograms[GreenImageChannel].data, 
									   cumulativeHistogram.intervalCount, g);

		b = GetEqualizedColorComponent(cumulativeHistogram.channelHistograms[BlueImageChannel].data, 
									   cumulativeHistogram.intervalCount, b);

		KPixel* newCurrentPixel = [[KPixel alloc] initWithRed:r green:g andBlue:b];
		
		[currentImage setPixel:newCurrentPixel atX:x y:y];
	}
}

- (void) equalizeBrightnessHistogram
{
	KCumulativeHistogram* cumulativeHistogram = currentImage.cumulativeHistogram;
	
	ForEachImagePixel(currentImage)
	{
		KPixel* currentPixel = [currentImage pixelAtX:x y:y];
		
		CGFloat brightness = [currentPixel brightness];
		CGFloat hue = [currentPixel hue];
		CGFloat saturation = [currentPixel saturation];
		
		brightness = 
		GetEqualizedColorComponent(
								cumulativeHistogram.channelHistograms[BrightnessImageChannel].data,
								cumulativeHistogram.intervalCount, brightness);
		
		
		KPixel* newCurrentPixel = 
		[[KPixel alloc] initWithHue:hue saturation:saturation andBrightness:brightness];
		
		[currentImage setPixel:newCurrentPixel atX:x y:y];
	}
}

- (void) applyCompositeEqualizationInternalWithLuminance: (CGFloat)luminanceFraction 
													 RGB: (CGFloat)rgbFraction 
													 Hue: (CGFloat)hueFraction
{
	KCumulativeHistogram* cumulativeHistogram = currentImage.cumulativeHistogram;
	
	ForEachImagePixel(currentImage)
	{
		KPixel* currentPixel = [currentImage pixelAtX:x y:y];
		CGFloat brightness = [currentPixel brightness];
		CGFloat hue = [currentPixel hue];
		CGFloat saturation = [currentPixel saturation];
		
		/* Luminance equalization */
		CGFloat eBrightness = 
		GetEqualizedColorComponent(
								cumulativeHistogram.channelHistograms[BrightnessImageChannel].data,
								currentImage.cumulativeHistogram.intervalCount,
								brightness);
		
		brightness = luminanceFraction * eBrightness + (1 - luminanceFraction) * brightness;
		
		currentPixel = 
		[[KPixel alloc] initWithHue:hue saturation:saturation andBrightness:brightness];

		/* RGB equalization */
		CGFloat r = [currentPixel red];
		CGFloat g = [currentPixel green];
		CGFloat b = [currentPixel blue];			
		
		CGFloat eR = 
		GetEqualizedColorComponent(cumulativeHistogram.channelHistograms[RedImageChannel].data,
								   currentImage.cumulativeHistogram.intervalCount, r);
		r = rgbFraction * eR + (1 - rgbFraction) * r;
		
		CGFloat eG = 
		GetEqualizedColorComponent(cumulativeHistogram.channelHistograms[GreenImageChannel].data,
								   currentImage.cumulativeHistogram.intervalCount, g);
		g = rgbFraction * eG + (1 - rgbFraction) * g;
		
		CGFloat eB = 
		GetEqualizedColorComponent(cumulativeHistogram.channelHistograms[BlueImageChannel].data,
								   currentImage.cumulativeHistogram.intervalCount, b);
		
		b = rgbFraction * eB + (1 - rgbFraction) * b;
		
		currentPixel = 
		[[KPixel alloc] initWithRed:r green:g andBlue:b];
		
		/* Hue equalization */
		brightness = [currentPixel brightness];
		hue = [currentPixel hue];
		saturation = [currentPixel saturation];
		
		CGFloat eHue = 
		GetEqualizedColorComponent(cumulativeHistogram.channelHistograms[HueImageChannel].data,
								   currentImage.cumulativeHistogram.intervalCount,
								   hue);
	
		hue = hueFraction * eHue + (1 - hueFraction) * hue;
		
		currentPixel = 
		[[KPixel alloc] initWithHue:hue saturation:saturation andBrightness:brightness];			
		
		[currentImage setPixel:currentPixel atX:x y:y];
	}
}

- (IBAction) applyEqualization: (id) sender
{
	[self backupCurrentImage];
	
	switch ([sender tag]) 
	{
		case 0:
			[self equalizeBrightnessHistogram];
			break;
		case 1:
			[self equalizeRGBHistogram];
			break;
		case 2:
			[self equalizeHueHistogram];
			break;
		default:
			break;
	}
	
	[currentImage.histogram computeForImage:currentImage];
	[currentImage.cumulativeHistogram computeForImage:currentImage];
	[self updateHistogramStatistics];
	[histogramView setNeedsDisplay:YES];
	[chromaticityHistogramView setNeedsDisplay:YES];
	[mainView setNeedsDisplay:YES];
}

- (IBAction) applyCompositeEqualization: (id) sender
{
	[compositeEqualizationWindow orderOut:self];
	[self backupCurrentImage];
	
	CGFloat luminanceFraction = [luminanceEqualizationSlider floatValue];
	CGFloat rgbFraction = [rgbEqualizationSlider floatValue];
	CGFloat hueFraction = [hueEqualizationSlider floatValue];
	
	[self applyCompositeEqualizationInternalWithLuminance:luminanceFraction 
													  RGB:rgbFraction 
													  Hue:hueFraction];
	
	[currentImage.histogram computeForImage:currentImage];
	[currentImage.cumulativeHistogram computeForImage:currentImage];
	[self updateHistogramStatistics];
	[histogramView setNeedsDisplay:YES];
	[chromaticityHistogramView setNeedsDisplay:YES];
	[mainView setNeedsDisplay:YES];
}

- (IBAction) showCompositeEqualizationWindow: (id) sender
{
	[compositeEqualizationWindow makeKeyAndOrderFront:compositeEqualizationWindow];
}

- (IBAction) showSegmentationWindow: (id) sender
{
	[channelForSegmentationComboBox selectItemAtIndex:[currentImage.histogram maxVarianceChannel]];
	[segmentationWindow makeKeyAndOrderFront:self];
}

- (IBAction) applyAutoLevels: (id) sender
{
	KHistogram* histogram = currentImage.histogram;
	NSUInteger* channelHistogramData = 
	histogram.channelHistograms[BrightnessImageChannel].data;
	CGFloat rejectionThresholdForShadows = 0.03;
	CGFloat rejectionThresholdForHighlights = 0.005;
	NSUInteger imagePixelCount = currentImage.width * currentImage.height;

	NSUInteger rejectedPixelCount = 0;
	NSUInteger blackLevelInterval = 0;
	while ((CGFloat)rejectedPixelCount / imagePixelCount < rejectionThresholdForShadows) 
	{
		rejectedPixelCount += channelHistogramData[blackLevelInterval];
		blackLevelInterval++;
	}	
	
	rejectedPixelCount = 0;
	NSUInteger whiteLevelInterval = histogram.intervalCount - 1;
	while ((CGFloat)rejectedPixelCount / imagePixelCount < rejectionThresholdForHighlights) 
	{
		rejectedPixelCount += channelHistogramData[whiteLevelInterval];
		whiteLevelInterval--;
	}
	
	CGFloat blackLevel = (CGFloat)blackLevelInterval / histogram.intervalCount;
	CGFloat whiteLevel = (CGFloat)whiteLevelInterval / histogram.intervalCount;
	CGFloat brightnessSpan = whiteLevel - blackLevel;
	
	[self backupCurrentImage];
	ForEachImagePixel(currentImage)
	{
		KPixel* currentPixel = [currentImage pixelAtX:x y:y];
		CGFloat brightness = [currentPixel brightness];
		CGFloat hue = [currentPixel hue];
		CGFloat saturation = [currentPixel saturation];
			
		brightness = (brightness - blackLevel) / brightnessSpan;
		
		KPixel* newCurrentPixel = 
		[[KPixel alloc] initWithHue:hue saturation:saturation andBrightness:brightness];
		
		[currentImage setPixel:newCurrentPixel atX:x y:y];
	}
	
	[currentImage.histogram computeForImage:currentImage];
	[currentImage.cumulativeHistogram computeForImage:currentImage];
	[self updateHistogramStatistics];
	[histogramView setNeedsDisplay:YES];
	[chromaticityHistogramView setNeedsDisplay:YES];
	[mainView setNeedsDisplay:YES];
}

- (IBAction) changeSegmentationSource: (id) sender
{
	[channelForSegmentationComboBox setEnabled:useChromaticityForSegmentation];
	
	useChromaticityForSegmentation = !useChromaticityForSegmentation;
}

- (IBAction) segmentImage: (id) sender
{
	[segmentationWindow orderOut:self];
/*	
	NSColor* rainbow[] = 
	{
		[NSColor brownColor],
		[NSColor blueColor],
		[NSColor cyanColor],
		[NSColor greenColor],
		[NSColor magentaColor]
	};
*/	
	NSUInteger segmentCount = [segmentCountTextField intValue];
	if (useChromaticityForSegmentation == NO)
	{		
		channelForSegmentation = [channelForSegmentationComboBox indexOfSelectedItem];
		
		if (segmentationMap != nil)
			[segmentationMap release];
		
		segmentationMap = [KImageSegmentator segmentImage:currentImage 
							  usingKMeansWithSegmentCount:segmentCount
										  andUsingChannel: channelForSegmentation];
	}
	else
	{
		if (chromaticitySegmentationMap != nil)
			[chromaticitySegmentationMap release];
		
		chromaticitySegmentationMap = [KImageSegmentator segmentImageByChromaticity:currentImage
														usingKMeansWithSegmentCount:segmentCount];
	}
/*	
	KImage* segmentationImage = 
	[[KImage alloc] initWithWidth:currentImage.width andHeight:currentImage.height];
	ForEachImagePixel(currentImage)
	{
		NSUInteger index = y * currentImage.width + x;
	}
	
	[self backupCurrentImage];
	displayBuffer = segmentationImage;
	[mainView setNeedsDisplay:YES];
 */
	[chromaticityHistogramView setNeedsDisplay:YES];
	[histogramView setNeedsDisplay:YES];
 
}

- (IBAction) showChromaticityHistogramWindow: (id) sender
{
	displayedChromaticityHistogramChannel = HueChromaticityChannel;
	[chromaticityHistogramChannelComboBox selectItemAtIndex:0];
	[chromaticityHistogramWindow makeKeyAndOrderFront:self];
}

- (IBAction) changeDisplayedChromaticityHistogramChannel: (id) sender
{
	displayedChromaticityHistogramChannel = 
	[chromaticityHistogramChannelComboBox indexOfSelectedItem];
	[chromaticityHistogramView setNeedsDisplay:YES];
}

- (void) dealloc
{
	[currentImage release];
	[segmentationMap release];
	
	[super dealloc];
}

@end
